#!/usr/bin/env python3
"""wikisplit.py - Split a wikipedia XML dump into
multiple files.

Copyright Matthew Leon Grinshpun, 2009
"""

import os
from optparse import OptionParser

endtag = "</mediawiki>"
article_endtag = "</page>"

def split_wikidump(dumpfile, splits):
    """Split dumpfile into [splits] files."""

    splitsize = os.path.getsize(dumpfile) // splits 
    with open(dumpfile) as input:
        wikitag = input.readline()
        output_files = (("{0}.split{1:=0{2}}").format(dumpfile, i, len(str(splits)))
                            for i in range(1, splits+1)) 
        for file in output_files:
            with open(file, 'w') as outfile:
                outfile.write(wikitag)
                outfile.writelines(input.readlines(splitsize))

                inline = True # prime the pump
                while inline:
                    inline = input.readline()
                    outfile.write(inline)
                    if article_endtag in inline:
                        break
                if not endtag in inline: # otherwise it's already written
                    outfile.write(endtag + "\n")

def main():
    parser = OptionParser(usage="usage: %prog [splits] [dumpfile]",
                          description="[splits] is an integer representing the\
                                       number of target files to generate.\
                                       \n[dumpfile] is a path to the XML \
                                       Wikipedia dump file.")
    try:
        (number_of_splits, dumpfile) = parser.parse_args()[1]
        number_of_splits = int(number_of_splits)
    except TypeError as e:
        print(e)
        parser.print_help()
        return 1

    split_wikidump(dumpfile, number_of_splits)

if __name__ == "__main__":
    exit(main())
